import pandas as pd
import os
import numpy as np


def load_data(raw_data_path, data_path):
    if not os.path.exists(data_path):
        os.makedirs(data_path)
    for path, _, filenames in os.walk(raw_data_path):
        for filename in filenames:
            data = pd.read_csv(os.path.join(path, filename))
            # 找到第一个敌机数据的索引
            if data[data['Id'] == '102'].empty == True or data[data['Id'] == '103'].empty == True:
                continue
            first_enemy_index = data[data['Id'] == '103'].index.values[0]
            data = data.loc[first_enemy_index - 1:]
            data.reset_index(drop=True, inplace=True)
            last_enemy_index = data.index.values[-1]
            for index in range(0, data.shape[0]):
                if data.loc[index, 'Id'] == '102' and data.loc[index + 1, 'Id'] != '103':
                    last_enemy_index = index - 1
                    break
                if data.loc[index, 'Id'] == '103' and data.loc[index - 1, 'Id'] != '102':
                    last_enemy_index = index - 1
                    break
            data = data.loc[:last_enemy_index]
            data.reset_index(drop=True, inplace=True)

            myPlane_index = data[data['Id'] == '102'].index.values
            for index in myPlane_index:
                data.loc[index, 'label'] = 0
                for missile_index in range(index + 2, last_enemy_index + 1):
                    if data.loc[missile_index, 'Id'] != '102' and data.loc[missile_index, 'Id'] != '103':
                        distance1 = (data.loc[index, 'Longitude'] - data.loc[missile_index, 'Longitude']) ** 2 + (
                                data.loc[index, 'Latitude'] - data.loc[missile_index, 'Latitude']) ** 2 + (
                                            data.loc[index, 'Altitude'] - data.loc[missile_index, 'Altitude']) ** 2
                        distance2 = (data.loc[index + 1, 'Longitude'] - data.loc[missile_index, 'Longitude']) ** 2 + (
                                data.loc[index + 1, 'Latitude'] - data.loc[missile_index, 'Latitude']) ** 2 + (
                                            data.loc[index + 1, 'Altitude'] - data.loc[missile_index, 'Altitude']) ** 2
                        if distance1 <= distance2:
                            data.loc[index, 'label'] = 1
                            break
                    else:
                        break
            myPlane = data[data['Id'] == '102'].copy()
            myPlane.reset_index(drop=True, inplace=True)
            enemy = data[data['Id'] == '103'].copy()
            enemy.reset_index(drop=True, inplace=True)
            for col in ['ISO time', 'Unix time', 'label']:
                del enemy[col]
            for col in enemy.columns:
                enemy.rename(columns={col: 'enemy_' + col}, inplace=True)
            df = pd.concat([myPlane, enemy], sort=False, axis=1)
            df.reset_index(drop=True, inplace=True)
            df = pd.concat(
                [df, pd.DataFrame({'V_x': df['U'].diff(), 'V_y': df['Altitude'].diff(), 'V_z': df['V'].diff()})],
                sort=False, axis=1)
            df.fillna(0, inplace=True)
            df[['V_x', 'V_y', 'V_z']] = df[['V_x', 'V_y', 'V_z']] / 0.05
            df['V_all'] = df.apply(lambda x: np.sqrt(x['V_x'] ** 2 + x['V_y'] ** 2 + x['V_z'] ** 2), axis=1)
            df['distance'] = df.apply(lambda x: np.sqrt(
                (x['U'] - x['enemy_U']) ** 2 + (x['Altitude'] - x['enemy_Altitude']) ** 2 + (
                        x['V'] - x['enemy_V']) ** 2), axis=1)
            df['cos'] = df.apply(lambda x: (x['V_x'] * (x['enemy_U'] - x['U']) + x['V_y'] * (
                        x['enemy_Altitude'] - x['Altitude']) + x['V_z'] * (x['enemy_V'] - x['V'])) / (
                                                       x['V_all'] * x['distance']) if x['V_all'] * x['distance'] != 0 else None, axis=1)
            df.reset_index(drop=True, inplace=True)
            df.to_csv(os.path.join(data_path, filename), index=False)
            print(filename + '已输出')


if __name__ == '__main__':
    load_data('raw_data', 'data')
    # print(data['Id'].head())
    # data.to_csv('1_new.csv', index=False)